"""
实现词向量和数字向量的转换
"""


class WordSequence:
    UNK_TAG = "UNK"  # 未定义的词
    PAD_TAG = "PAD"  # 填充词

    UNK = 0
    PAD = 1

    def __init__(self):
        # 定义词典
        self.dict = {
            self.UNK_TAG: self.UNK,
            self.PAD_TAG: self.PAD
        }
        # 定义统计词频的字典
        self.count_dict = {}
        # 定义反转字典
        self.reverse_dict = {}

    def fit(self, sentences):
        """
        统计词频
        :param sentences: [word1, word2, ...]
        :return:
        """
        for word in sentences:
            self.count_dict[word] = self.count_dict.get(word, 0) + 1

    def build_dict(self, min_rate=None, max_rate=None, max_word=None):
        """
        创建字典
        :param min_rate: 最小词频
        :param max_rate: 最大词频
        :param max_word: 字典最大长度
        :return:
        """
        # 取最小词频要求的词
        if min_rate:
            self.count_dict = {word: value for word, value in self.count_dict.items() if value >= min_rate}

        # 取最大词频要求的词
        if max_rate:
            self.count_dict = {word: value for word, value in self.count_dict.items() if value <= max_rate}

        # 创建词典
        if max_word and (len(self.count_dict) > max_word):
            self.count_dict = dict(sorted(list(self.count_dict.items()), key=lambda x: x[1], reverse=True)[:max_word])
        for word in self.count_dict.keys():
            self.dict[word] = len(self.dict)

        # 创建反转字典
        self.reverse_dict = dict(zip(self.dict.values(), self.dict.keys()))

    def word_to_num_transform(self, word_vector, max_len=None):
        """
        将词向量转换为数字向量
        :param word_vector: [word1, word2, ...]
        :param max_len: 最大长度
        :return:
        """
        # 1. 遍历词向量，查询每个词在词语字典中对应的数字，创建数字向量
        num_vector = [self.dict.get(word, self.UNK) for word in word_vector]

        # 2. 根据最大长度要求修改数字向量
        if max_len:
            if len(num_vector) > max_len:
                num_vector = num_vector[:max_len]
            else:
                num_vector.extend([self.PAD] * (max_len - len(num_vector)))

        # 3. 返回数字向量
        return num_vector

    def num_to_word_transform(self, num_vector, max_len=None):
        """
        将数字向量转换为词向量
        :param num_vector: [num1, num2, ...]
        :param max_len: 最大长度
        :return:
        """
        # 1. 遍历数字向量， 查询每个数字在反转字典中对应的词，创建词向量
        word_vector = [self.reverse_dict.get(num, self.UNK_TAG) for num in num_vector]

        # 2. 根据最大长度要求修改词向量
        if max_len:
            if len(word_vector) > max_len:
                word_vector = word_vector[:max_len]
            else:
                word_vector.extend([self.PAD_TAG] * (max_len - len(word_vector)))

        # 3. 返回词向量
        return word_vector

    def __len__(self):
        return len(self.dict)


if __name__ == '__main__':
    ws = WordSequence()
    ws.fit(["我", "是", "谁"])
    ws.fit(["我", "在", "哪"])
    ws.build_dict()

    print(ws.word_to_num_transform(["你", "我", "他", "她"], max_len=20))

    print(ws.num_to_word_transform([1, 2, 4, 5, 7, 2, 3], max_len=20))

    print(ws.dict)
    print(ws.reverse_dict)
